import argparse
import torch

from torchvision.datasets import MNIST
from torch.utils.data import DataLoader
from lenet.config import DEFAULT_TRANSFORMS
from lenet.model import LeNet


def test(args):
    # loading the dataset and preprocessing
    test_dataset = MNIST(
        root="./data",
        train=False,
        transform=DEFAULT_TRANSFORMS,
        download=True,
    )
    test_loader = DataLoader(
        dataset=test_dataset,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        shuffle=True,
    )

    model = LeNet().to(args.device)
    if args.weights_file:
        model.load_state_dict(torch.load(args.weights_file, map_location=args.device))

    with torch.no_grad():
        correct = 0
        total = 0
        for images, labels in test_loader:
            images = images.to(args.device)
            labels = labels.to(args.device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        print(
            "Accuracy of the network on the {} test images: {} %".format(
                total, 100 * correct / total
            )
        )


def parse_args():
    """
    Parse input arguments
    """
    parser = argparse.ArgumentParser(description="LeNet training")

    # fmt: off
    parser.add_argument('--device', default='cuda:0', type=str)
    parser.add_argument('--batch-size', default=64, type=int, help='batch size for data loader')
    parser.add_argument('--num-workers', default=4, type=int, help='number of workers for data loader')
    parser.add_argument("--weights-file", default="output/lenet_ckpt_100.pth", type=str, help="weights file")
    # fmt: on

    args = parser.parse_args()
    return args


if __name__ == "__main__":
    args = parse_args()
    test(args)
